from dataset import Dataset
from data_generator import DataGenerator
from numpy.random import RandomState
import numpy as np
from preloaded_map import preloaded_map
from numbers import Integral, Real
from collections import Callable
from dataset import map


class _mapper(Callable):
    def __init__(self, mapped):
        assert isinstance(mapped, Callable)
        self._mapped = mapped

    def __call__(self, dataset, *args, **kwargs):
        assert isinstance(dataset, Dataset)
        return map(self._mapped(*args, **kwargs), dataset)


def _flip(x_axes):
    x, axes = x_axes
    # axes: bool vector indicating which axes to perform the flip
    slices = tuple(slice(None, None, s) for s in (-1 if axis else 1 for axis in axes))
    return x[slices]


def random_flip(data_generator, prob, rng, preloaded_len, **pool_kwargs):
    assert isinstance(data_generator, DataGenerator)
    assert isinstance(rng, RandomState)
    prob = np.asarray(prob, dtype='float')
    assert prob.ndim <= 1

    def gen():
        while True:
            x = next(data_generator)
            x = np.asarray(x)
            axes = rng.rand(x.ndim) < prob
            yield x, axes

    return preloaded_map(_flip, gen(), preloaded_len=preloaded_len, **pool_kwargs)


def _crop(x_starts_shape):
    x, starts, shape = x_starts_shape
    slices = tuple(slice(st, st+sp) for st, sp in zip(starts, shape))
    croped = x[slices]
    assert croped.shape == tuple(shape)
    return croped


def random_crop(data_generator, target_shape, rng, preloaded_len, **pool_kwargs):
    assert isinstance(data_generator, DataGenerator)
    target_shape = np.asarray(target_shape)
    assert all(isinstance(ts, Integral) and ts > 0 for ts in target_shape)
    assert isinstance(rng, RandomState)

    def gen():
        while True:
            x = next(data_generator)
            x = np.asarray(x)
            diff = np.asarray(x.shape) - target_shape
            starts = tuple(rng.randint(s + 1) for s in diff)
            yield x, starts, target_shape

    return preloaded_map(_crop, gen(), preloaded_len=preloaded_len, **pool_kwargs)


class _parameterized_crop(Callable):
    def __init__(self, target_shape):
        self._target_shape = np.asarray(target_shape)
        assert all(isinstance(ts, Integral) and ts > 0 for ts in self._target_shape)

    def __call__(self, x_starts):
        x, starts = x_starts
        x = np.asarray(x)
        croped = _crop((x, starts, self._target_shape))
        return croped


def parameterized_crop(dataset, target_shape):
    assert isinstance(dataset, Dataset)
    return map(_parameterized_crop(target_shape), dataset)


class _middle_crop(Callable):
    def __init__(self, target_shape):
        self._target_shape = np.asarray(target_shape)
        assert all(isinstance(ts, Integral) and ts > 0 for ts in self._target_shape)

    def __call__(self, x):
        x = np.asarray(x)
        diff = np.asarray(x.shape) - self._target_shape
        starts = diff // 2
        croped = _crop((x, starts, self._target_shape))
        return croped


def middle_crop(dataset, target_shape):
    assert isinstance(dataset, Dataset)
    return map(_middle_crop(target_shape), dataset)


class _global_normalization(Callable):
    def __init__(self, mean, std):
        assert isinstance(mean, Real)
        assert isinstance(std, Real)
        self._mean = mean
        self._std = std

    def __call__(self, x):
        x = np.asarray(x)
        normalized = (x - self._mean) / self._std
        return normalized


def global_normalization(dataset, mean, std):
    assert isinstance(dataset, Dataset)
    return map(_global_normalization(mean, std), dataset)


class _astype(Callable):
    def __init__(self, dtype):
        self._dtype = dtype

    def __call__(self, x):
        typed = np.asarray(x, dtype=self._dtype)
        return typed


def astype(dataset, dtype):
    assert isinstance(dataset, Dataset)
    return map(_astype(dtype), dataset)


class _append_channel(Callable):
    def __call__(self, x):
        x = np.asarray(x)
        return x[..., None]


def append_channel(dataset):
    assert isinstance(dataset, Dataset)
    return map(_append_channel(), dataset)


if __name__ == '__main__':
    in_data = np.arange(25).reshape((5, 5))
    in_data = np.repeat(in_data[None], 2, axis=0)
    print(in_data)
    from dataset import product
    in_data = product(in_data, [(0,0), (2,2), (1,1)])
    out_data = parameterized_crop(in_data, (3, 3))
    print(len(out_data))
    for i in range(len(out_data)):
        print(out_data[i])
    import pickle
    s = pickle.dumps(out_data, -1)
